package spark.example

import org.apache.spark.{SparkConf,SparkContext}
import org.apache.spark.ml.feature.Word2Vec
import org.apache.spark.ml.linalg.Vector
import org.apache.spark.sql.Row
import org.apache.spark.sql.types._

object word2Vec {
    def main(args: Array[String]) {
    val sparkConf = new SparkConf().setAppName("word2Vec")
    val sc = new SparkContext(sparkConf)
    val spark = new org.apache.spark.sql.SQLContext(sc)
    val documentDF = spark.createDataFrame(Seq(
        "Hi I heard about Spark".split(" "),
        "I wish Java could use case classes".split(" "),
        "Logistic regression modles are neat".split(" ")
    ).map(Tuple1.apply)).toDF("text")

    val word2Vec = new Word2Vec()
        .setInputCol("text")
        .setOutputCol("result")
        .setVectorSize(3)
        .setMinCount(0)
        
    val model = word2Vec.fit(documentDF)

    val result = model.transform(documentDF)
    result.collect().foreach { case Row(text:Seq[_],features:Vector) =>
        println(s"Text:[${text.mkString(", v")}] => \nVector: $features\n") }
    }
}
